"""argparser configuration"""

import argparse


def add_prompt_args(parser):
    """prompt arguments"""

    group = parser.add_argument_group('Prompt', 'Prompt configurations')
    group.add_argument('--prompt_template', type=str, default="prompt.txt",
                       help="path to the original prompt template file")
    group.add_argument('--input_keys', type=str, nargs='+', help="input keys")
    group.add_argument('--label_key', type=str, default="label", help="label key")
    group.add_argument('--label_type', type=str, default="str", help="label type")
    return parser


def add_config_args(parser):
    """config arguments"""

    group = parser.add_argument_group('Settings', 'General settings')
    group.add_argument('--threads', type=int, default=20, help="thread number")
    group.add_argument('--seed', type=int, default=42, help="random seed")
    group.add_argument('--model', type=str, default="gpt-3.5-turbo", help="LLM model for evaluation")
    group.add_argument('--save', type=str, default="prompt.opt.txt",
                       help="file to save the optimised prompt template")

    return parser


def add_ga_args(parser):
    """GA arguments"""

    group = parser.add_argument_group('Genetic Algorithm', 'GA configurations')

    group.add_argument('--population_size', type=int, default=64,
                       help="number of individuals in each population")

    group.add_argument('--max_generations', type=int, default=5,
                       help="number of iterations")

    group.add_argument('--mutation_rate', type=float, default=0.75,
                       help="gene mutation probability")

    group.add_argument('--crossover_rate', type=float, default=0.5,
                       help="gene crossover probability")

    group.add_argument('--deletion_prob', type=float, default=0.1,
                       help="delete probability in mutation")

    group.add_argument('--paraphrasing_prob', type=float, default=0.3,
                       help="paraphrasing probability in mutation")

    group.add_argument('--translation_prob', type=float, default=0.3,
                          help="translation probability in mutation")

    group.add_argument('--expansion_prob', type=float, default=0.2,
                       help="expansion probability in mutation")

    group.add_argument('--exchanging_prob', type=float, default=0.1,
                            help="exchanging probability in mutation")


    return parser


def add_data_args(parser):
    """Train/valid/test data arguments."""

    group = parser.add_argument_group('Data', 'Data configurations')

    group.add_argument('--k_shot', type=int, default=20,
                       help="number of training examples for each label class")

    group.add_argument('--data_path', type=str, default="data.jsonl",
                       help="path to data")

    return parser


def add_tool_args(parser):
    """tool arguments"""

    group = parser.add_argument_group('Tool', 'Tool configurations')
    group.add_argument('--paraphraser_prompt', type=str, default="tool_prompts/paraphraser.txt",
                          help="prompt for paraphraser")
    group.add_argument('--paraphraser_model', type=str, default="gpt-3.5-turbo-instruct", #"gpt-4-1106-preview",
                            help="paraphraser model")
    group.add_argument('--translator_prompt', type=str, default="tool_prompts/translator.txt",
                       help="prompt for translator")
    group.add_argument('--translator_model', type=str, default="gpt-3.5-turbo-instruct",
                       help="translator model")

    return parser


def get_args():
    """Parse all the args."""

    parser = argparse.ArgumentParser(description='GA Prompt Optimisation')
    parser = add_ga_args(parser)
    parser = add_config_args(parser)
    parser = add_data_args(parser)
    parser = add_prompt_args(parser)
    parser = add_tool_args(parser)

    args = parser.parse_args()

    return args
